#!/bin/bash
# qwen3-embedding-safe-upgrade.sh
# 安全升级 Qwen3-Embedding 到 Q5_K_M
# 步骤：下载 → 构建新模型 → 验证 → 手动替换

set -e

# =============== 配置区 ===============
MODEL_URL="https://hf-mirror.com/Qwen/Qwen3-Embedding-8B-GGUF/resolve/main/Qwen3-Embedding-8B-Q5_K_M.gguf"
MODEL_FILE="Qwen3-Embedding-8B-Q5_K_M.gguf"
TEMP_DIR="/tmp/ollama-qwen3"
NEW_MODEL_NAME="qwen3-embedding:8b-q5km"    # 新名字
OLD_MODEL_NAME="qwen3-embedding:8b"         # 旧名字
# =====================================

echo "🚀 开始安全升级 Qwen3-Embedding 到 Q5_K_M"

# 1. 创建临时目录
mkdir -p $TEMP_DIR
cd $TEMP_DIR

# 2. 下载模型
# rm -f "$MODEL_FILE"
if [ ! -f "$MODEL_FILE" ]; then
    echo "⬇️ 下载模型: $MODEL_URL"
    # wget -O "$MODEL_FILE" "$MODEL_URL"
    aria2c -x 16 -s 16 -k 16M -o "$MODEL_FILE" "$MODEL_URL"
else
    echo "✅ 模型已存在，跳过下载"
fi

# 3. 创建 Modelfile（使用相对路径！）
cat > Modelfile << 'EOF'
# Modelfile - 使用本地文件（相对路径）
FROM ./Qwen3-Embedding-8B-Q5_K_M.gguf

# ✅ 标准参数
PARAMETER num_ctx 8192
PARAMETER num_gpu 99

TEMPLATE """{{- if .Messages }}
{{- range $i, $_ := .Messages }}
{{- if eq $i 0 }}
<|system|>
You are a helpful assistant.
Embedding only.
{{- end }}
{{- end }}
{{- end }}
{{ .Prompt }}"""
EOF

echo "✅ Modelfile 已创建"

# 4. 构建新模型（Ollama 会自动复制文件到 blobs）
echo "🔨 构建新模型: $NEW_MODEL_NAME"
ollama create "$NEW_MODEL_NAME" -f Modelfile

# 5. 验证新模型
echo "🔍 验证新模型是否可用..."
echo "📌 请手动运行测试："
echo "   ollama run $NEW_MODEL_NAME"
echo "   或测试 API："
echo "   curl http://localhost:11434/api/embeddings -d '{\"model\":\"$NEW_MODEL_NAME\", \"prompt\":\"test\"}'"

# 6. 显示当前模型列表
echo
echo "📋 当前模型列表："
ollama list

# 7. 提示用户手动替换
echo
echo "✅ 新模型已就绪！"
echo "💡 建议操作流程："
echo "   1. 在 RAGFlow / 应用中先测试 $NEW_MODEL_NAME"
echo "   2. 确认效果满意后，再执行："
echo "      ollama rm $OLD_MODEL_NAME"
echo "      ollama tag $NEW_MODEL_NAME $OLD_MODEL_NAME"
echo "   3. 重启应用，即可无缝切换"

echo "🎉 升级完成！安全、可回滚、零风险。"